package com.llmops.demo;

import com.alibaba.fastjson.JSON;
import com.llmops.core.*;
import com.llmops.core.node.LLMNode;
import com.llmops.core.node.StartNode;
import com.llmops.core.workflow.WorkflowParser;

import java.util.ArrayList;
import java.util.List;

public class GraphConfigTest {

    public static void main(String[] args) {
//        WorkflowGraph workflow = new WorkflowGraph();
//        Context context = new Context();
//        context.setVariable("sessionId",1);
//
//        workflow.addNode(new UserInputNode("data_input","提示词"));
//        workflow.addNode(new ClassificationNode("llm_ClassificationNode", Map.of("classifications",List.of("售前","售后"))));
////        workflow.addNode(new LLMNode("llm_inferenceB", Collections.singletonMap("model", "gpt-B")));
//
//
//        workflow.addNode(new OutputNode("result_exporterA", Collections.singletonMap("v", "A")));
//        workflow.addNode(new OutputNode("result_exporterB", Collections.singletonMap("v", "B")));
//
//        workflow.addEdge(new Edge("data_input", "llm_ClassificationNode"));
////        workflow.addEdge(new Edge("data_loader", "llm_inferenceB", Edge.ConditionType.EQUAL,"B"));
//
//        workflow.addEdge(new Edge("llm_ClassificationNode", "result_exporterA", Edge.ConditionType.EQUAL,"售前"));
//        workflow.addEdge(new Edge("llm_ClassificationNode", "result_exporterB", Edge.ConditionType.EQUAL,"售后"));
//
////        workflow.addEdge(new Edge("llm_inferenceB", "result_exporterB"));
//
//        workflow.validate();
//        workflow.execute(context);
//
////        workflow.execute(context);
//
//        System.out.println("Workflow execution completed");

//        getConfig();
        getWg();
    }

    private static GraphConfig getConfig(){
        List<NodeProperty> nodes = new ArrayList<>();
        List<Edge> edges = new ArrayList<>();
        GraphConfig config = new GraphConfig();

        StartNode.UserInputNodeProperty config0 = new StartNode.UserInputNodeProperty();
        config0.setId("input_1");
        config0.setType(StartNode.class.getSimpleName());
        config0.setInput("你是谁");
        nodes.add(config0);

        LLMNode.LLMNodeProperty config1 = new LLMNode.LLMNodeProperty();
        config1.setId("llm_1");
        config1.setType(LLMNode.class.getSimpleName());
        config1.setApiUrl("123");
        nodes.add(config1);

        config.setNodes(nodes);

        Edge edge = new Edge("input_1","llm_1");
        edges.add(edge);
        config.setEdges(edges);

        System.out.printf(JSON.toJSONString(config));
        return config;
    }

    static WorkflowGraph getWg(){
        Context context = new Context();
        WorkflowParser parser= new WorkflowParser();
        String json="{\"connections\":[{\"sourceId\":\"input_1\",\"targetId\":\"llm_1\"}],\"nodes\":[{\"id\":\"input_1\",\"config\":{\"input\":\"你是谁\"},\"type\":\"userInputNode\"},{\"config\":{\"prompt\":\"你是聊天助手请回答我的问题：{input_1_output}\"},\"id\":\"llm_1\",\"type\":\"llmNode\"}],\"id\":\"1\"}";
        WorkflowGraph workflowGraph = parser.parserFlow(json);
        workflowGraph.execute(context);

        String s = parser.parserFlow(workflowGraph);
        System.out.println(s);
        return null;
    }
}
